// pc_udp_sender.cpp — RAW PointCloud2 passthrough with optional Zstd
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <ros/serialization.h>

#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <endian.h>
#include <zstd.h>   // <— Zstd

#include <vector>
#include <string>
#include <cstring>
#include <algorithm>

#pragma pack(push, 1)
struct UdpHeader {
  uint32_t magic;         // 'PCUD' 0x50435544 (network order)
  uint32_t seq;           // frame sequence (network order)
  uint16_t total_chunks;  // total datagrams in this frame (network order)
  uint16_t idx;           // datagram index (network order)
  uint32_t payload_size;  // bytes after header (network order)
  uint64_t stamp_ns;      // original msg stamp in ns (big-endian)
  uint8_t  codec;         // 0=raw, 1=zstd (no byte order needed)
  uint8_t  reserved[3];   // padding/alignment
};
#pragma pack(pop)

class PcSender {
public:
  explicit PcSender(ros::NodeHandle& nh_priv) : nh_priv_(nh_priv) {
    nh_priv_.param<std::string>("in_cloud_topic", in_topic_, std::string("/cloud_registered"));
    nh_priv_.param<std::string>("dest_ip", dest_ip_, std::string("192.168.1.20"));
    nh_priv_.param<int>("dest_port", dest_port_, 50000);
    nh_priv_.param<int>("mtu", mtu_, 1300);
    nh_priv_.param<int>("queue_size", queue_size_, 1);
    nh_priv_.param<bool>("enable_stats", enable_stats_, true);

    // Zstd 开关与等级（1..5 建议; >6 压得更小但耗时上升）
    nh_priv_.param<bool>("use_zstd", use_zstd_, true);
    nh_priv_.param<int>("zstd_level", zstd_level_, 1);

    sock_ = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_ < 0) { perror("socket"); throw std::runtime_error("socket() failed"); }
    int sndbuf = 32 * 1024 * 1024;
    setsockopt(sock_, SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf));

    std::memset(&addr_, 0, sizeof(addr_));
    addr_.sin_family = AF_INET;
    addr_.sin_port   = htons(dest_port_);
    if (inet_pton(AF_INET, dest_ip_.c_str(), &addr_.sin_addr) != 1)
      throw std::runtime_error("invalid dest_ip");

    sub_ = nh_priv_.subscribe(in_topic_, queue_size_, &PcSender::cb, this,
                              ros::TransportHints().tcpNoDelay());

    ROS_INFO_STREAM("pc_udp_sender (RAW"
                    << (use_zstd_ ? "+Zstd" : "")
                    << ") -> " << dest_ip_ << ":" << dest_port_
                    << " topic=" << in_topic_ << " mtu=" << mtu_
                    << (use_zstd_ ? (" zstd_level=" + std::to_string(zstd_level_)) : ""));
  }

  ~PcSender(){ if (sock_ >= 0) close(sock_); }

private:
  void cb(const sensor_msgs::PointCloud2ConstPtr& msg) {
    static uint32_t seq = 0;

    // 1) 序列化（无损）
    const size_t raw_len = ros::serialization::serializationLength(*msg);
    std::vector<uint8_t> raw(raw_len);
    {
      ros::serialization::OStream os(raw.data(), raw.size());
      ros::serialization::serialize(os, *msg);
    }

    // 2) 可选 Zstd 压缩
    std::vector<uint8_t> payload;
    uint8_t codec = 0; // 0=raw,1=zstd
    if (use_zstd_) {
      size_t bound = ZSTD_compressBound(raw.size());
      payload.resize(bound);
      size_t compSize = ZSTD_compress(payload.data(), bound, raw.data(), raw.size(), zstd_level_);
      if (!ZSTD_isError(compSize) && compSize > 0 && compSize < raw.size()) {
        payload.resize(compSize);
        codec = 1;
      } else {
        // 压缩没收益或失败，就发原始
        payload.swap(raw); // payload=raw
        codec = 0;
        if (ZSTD_isError(compSize)) {
          ROS_WARN_THROTTLE(1.0, "zstd compress error: %s", ZSTD_getErrorName(compSize));
        }
      }
    } else {
      payload.swap(raw);
    }

    // 3) UDP 分片
    const size_t header_sz  = sizeof(UdpHeader);
    const size_t payload_mt = static_cast<size_t>(std::max(100, mtu_ - static_cast<int>(header_sz)));
    const uint16_t total    = static_cast<uint16_t>((payload.size() + payload_mt - 1) / payload_mt);

    UdpHeader h{};
    h.magic        = htonl(0x50435544);
    h.seq          = htonl(seq);
    h.total_chunks = htons(total);
    h.stamp_ns     = htobe64(static_cast<uint64_t>(msg->header.stamp.toNSec()));
    h.codec        = codec;

    size_t sent_chunks = 0;
    for (uint16_t i = 0; i < total; ++i) {
      const size_t off = static_cast<size_t>(i) * payload_mt;
      const size_t n   = std::min(payload_mt, payload.size() - off);
      h.idx            = htons(i);
      h.payload_size   = htonl(static_cast<uint32_t>(n));

      std::vector<uint8_t> buf(header_sz + n);
      std::memcpy(buf.data(), &h, header_sz);
      std::memcpy(buf.data() + header_sz, payload.data() + off, n);

      const ssize_t sent = sendto(sock_, buf.data(), buf.size(), 0,
                                  reinterpret_cast<sockaddr*>(&addr_), sizeof(addr_));
      if (sent > 0) ++sent_chunks;
    }

    if (enable_stats_) {
      static double last = 0;
      double now = ros::Time::now().toSec();
      if (now - last > 1.0) {
        ROS_INFO("[TX] seq=%u chunks=%u/%u bytes=%zu (codec=%s)",
                 seq,
                 static_cast<unsigned>(sent_chunks), static_cast<unsigned>(total),
                 payload.size(), codec ? "zstd" : "raw");
        last = now;
      }
    }
    ++seq;
  }

  ros::NodeHandle nh_priv_;
  ros::Subscriber sub_;
  std::string in_topic_, dest_ip_;
  int dest_port_{50000}, mtu_{1300}, queue_size_{1};
  bool enable_stats_{true};

  bool use_zstd_{true};
  int  zstd_level_{1};

  int sock_{-1};
  sockaddr_in addr_{};
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "pc_udp_sender");
  ros::NodeHandle nh_priv("~");
  try { PcSender s(nh_priv); ros::spin(); }
  catch (const std::exception& e) { ROS_FATAL("sender init failed: %s", e.what()); return 1; }
  return 0;
}

